#include <device/block.h>
#include <fflags.h>

#include <mm/slab.h>
#include <mm/valloc.h>

#include <lib/vsprintf.h>
#include <lib/math.h>

#define LOG_TAG "BLOCK"
#include <log.h>

static kmem_cache_t* blk_kmem_cache = NULL;

static lnode_t blk_devs;

static u32 blk_dev_count = 0;

void block_init() {
    blk_kmem_cache = new_kmem_cache("block_device", sizeof(blk_dev_t), 1, 0);
    list_head_init(&blk_devs);
}

static int blk_read_amount(blk_dev_t* bdev, void* buf, u64 start, size_t count) {
    pbuffer_t* pb = NULL;
    pbuffer_add(&pb, buf, bdev->blk_size * count);

    blkio_req_t* req = blkio_read_req(pb, start, 0);
    blkio_commit(bdev->blkio, req, BLKIO_WAIT);

    int errno = req->errcode;
    if (!errno) {
        errno = count;
    } else {
        errno = -errno;
    }

    blkio_free_req(req);
    pbuffer_free(pb);

    return errno;
}

static int blk_write_amount(blk_dev_t* bdev, void* buf, u64 start, size_t count) {
    pbuffer_t* pb = NULL;
    pbuffer_add(&pb, buf, bdev->blk_size * count);

    blkio_req_t* req = blkio_write_req(pb, start, 0);
    blkio_commit(bdev->blkio, req, BLKIO_WAIT);

    int errno = req->errcode;
    if (!errno) {
        errno = count;
    } else {
        errno = -errno;
    }

    blkio_free_req(req);
    pbuffer_free(pb);

    return errno;
}

static int blk_read(device_t* dev, void* buf, size_t offset, size_t len) {
    int errno;
    blk_dev_t* bdev = (blk_dev_t*)dev->impl;
    size_t bsize = bdev->blk_size;
    size_t rd_block = offset / bsize + bdev->start_lba;
    size_t r = offset % bsize;
    size_t rd_size = 0;

    if (!(len = MIN(len, ((size_t)bdev->end_lba - rd_block + 1) * bsize))) {
        return 0;
    }

    pbuffer_t* pb = NULL;
    blkio_req_t* req;
    void* head_buf = NULL, * tail_buf = NULL;

    // align the boundary
    if (r || len < bsize) {
        head_buf = valloc(bsize);
        rd_size = MIN(len, bsize - r);
        pbuffer_add(&pb, head_buf, bsize);
    }

    // align the length
    if ((len - rd_size)) {
        size_t llen = len - rd_size;
        assert_msg(!(llen % bsize), "misalign block read");
        pbuffer_add(&pb, buf + rd_size, llen);
    }

    req = blkio_read_req(pb, rd_block, 0);
    blkio_commit(bdev->blkio, req, BLKIO_WAIT);

    if (!(errno = req->errcode)) {
        memcpy(buf, head_buf + r, rd_size);
        errno = len;
    } else {
        errno = -errno;
    }

    if (head_buf) {
        vfree(head_buf);
    }

    blkio_free_req(req);
    pbuffer_free(pb);
    return errno;
}

static int blk_write(device_t* dev, void* buf, size_t offset, size_t len) {
    blk_dev_t* bdev = (blk_dev_t*)dev->impl;
    size_t bsize = bdev->blk_size;
    size_t wr_block = offset / bsize + bdev->start_lba;
    size_t r = offset % bsize;
    size_t wr_size = 0;

    if (!(len = MIN(len, ((size_t)bdev->end_lba - wr_block + 1) * bsize))) {
        return 0;
    }

    pbuffer_t* pb = NULL;
    blkio_req_t* req;
    void* tmp_buf = NULL;

    if (r) {
        size_t wr_size = MIN(len, bsize - r);
        tmp_buf = vzalloc(bsize);
        pbuffer_add(&pb, tmp_buf, bsize);

        memcpy(tmp_buf + r, buf, wr_size);
    }

    if ((len - wr_size)) {
        size_t llen = len - wr_size;
        assert_msg(!(llen % bsize), "misalign block write");
        pbuffer_add(&pb, buf + wr_size, llen);
    }

    req = blkio_write_req(pb, wr_block, 0);
    blkio_commit(bdev->blkio, req, BLKIO_WAIT);

    int errno = req->errcode;
    if (!errno) {
        errno = len;
    } else {
        errno = -errno;
    }

    if (tmp_buf) {
        vfree(tmp_buf);
    }

    blkio_free_req(req);
    pbuffer_free(pb);
    return errno;
}

static void blk_register(blk_dev_t* bdev) {
    char name[DEV_ID_SIZE + 1];
    int len = sprintf(name, "sd%c", 'a' + blk_dev_count++);
    device_t* dev = device_add(NULL, bdev, name, len, FS_VOLUME);
    dev->write = blk_write;
    dev->read = blk_read;

    bdev->dev = dev;
    strncpy(bdev->bdev_id, dev->name.val, DEV_ID_SIZE);
    list_append(&blk_devs, &bdev->list);
}

int block_mount(blk_dev_t* bdev) {

    blk_register(bdev);

    int errno = blkpart_probe(bdev->dev);
    if (errno < 0) {
        LOGE("Fail to parse partition on: %s (%d)\n", bdev->name, errno);
        return errno;
    }

    return errno;
}

blk_dev_t* blk_mount_part(blk_dev_t* bdev, const char* name, size_t index, u64 start_lba, u64 end_lba) {
    blk_dev_t* pbdev = take_obj(blk_kmem_cache);
    memcpy(pbdev, bdev, sizeof(*bdev));

    char part_name[PARTITION_NAME_SIZE + 1];

    int len = sprintf(part_name, "%sp%d", bdev->bdev_id, index);
    device_t* dev = device_add(NULL, pbdev, part_name, len, FS_VOLUME);
    dev->write = blk_write;
    dev->read = blk_read;

    pbdev->start_lba = start_lba;
    pbdev->end_lba = end_lba;
    pbdev->dev = dev;

    strcpy(pbdev->bdev_id, dev->name.val);
    if (name) {
        strncpy(pbdev->name, name, PARTITION_NAME_SIZE);
    }

    list_append(&bdev->parts, &pbdev->parts);
    return pbdev;
}

blk_dev_t* block_alloc_dev(const char* blk_id, void* driver, io_req_handler ioreq_handler) {
    blk_dev_t* bdev = take_obj(blk_kmem_cache);
    memset(bdev, 0, sizeof(blk_dev_t));
    list_head_init(&bdev->parts);
    strncpy(bdev->name, blk_id, PARTITION_NAME_SIZE);

    bdev->blkio = blkio_newctx(ioreq_handler);
    bdev->blkio->driver = driver;

    bdev->block_read = blk_read_amount;
    bdev->block_write = blk_write_amount;

    return bdev;
}